import cv2
import os
import time
import ril_env
import pygame
from ril.util import extract_min_rect, show_obj
from ril_env.envs import FlappyBirdEnv, Pipe, Bird
from ril_env import env_factory

BIRD_X = [72, 76]
BIRD_S = [32*22, 36*26]
UPPER_PIPE_W = [50, 54]

prev_bird_y, bird_v, prev_upper_pipe_y = None, 0, None


def flappy_bird_rule(frame, act: int):
    global prev_bird_y, bird_v, prev_upper_pipe_y
    if frame is None or frame.shape == (4, 84, 84):
        return act, None, True
    img = cv2.transpose(frame)
    objects = extract_min_rect(img, blur=None)  # 查找物体边缘，得到每个物体的最小矩形
    bird_y, upper_pipe_y = None, None
    # cv2.destroyAllWindows()
    for (index, (x, y, w, h)) in enumerate(objects):
        # if w > 36:
        #     show_obj(img, (x + w/2, y + h/2), w, h, str(index))
        if BIRD_X[0] < x + w / 2 < BIRD_X[1] and BIRD_S[0] < w*h < BIRD_S[1]:
            bird_y = y
        if UPPER_PIPE_W[0] < w < UPPER_PIPE_W[1]:
            upper_pipe_y = y + h

    if bird_y is None:
        if prev_bird_y is None:
            return FlappyBirdEnv.DO_NOTHING, None, False
        prev_bird_y += bird_v
        bird_y = prev_bird_y  # 如果无法识别到鸟，就用上一次的 y
    else:
        # print('update bird')
        prev_bird_y = bird_y
    if upper_pipe_y is None:
        upper_pipe_y = prev_upper_pipe_y
    else:
        prev_upper_pipe_y = upper_pipe_y

    act = FlappyBirdEnv.DO_NOTHING
    if bird_y is None or upper_pipe_y is None:
        act = FlappyBirdEnv.rand_action()
    elif bird_y >= upper_pipe_y + Pipe.GAP_SIZE/2 + Bird.FLAP_VEL + 6:
        act = FlappyBirdEnv.FLAP
    bird_v = Bird.FLAP_VEL if act is FlappyBirdEnv.FLAP else bird_v + Bird.ACC_Y
    # if upper_pipe_y is not None:
    #     print(bird_y, upper_pipe_y + Pipe.GAP_SIZE/2 + Bird.FLAP_VEL)
    return act, None, False


if __name__ == '__main__':
    try:
        del os.environ['SDL_VIDEODRIVER']
    except KeyError:
        pass
    env = env_factory('FlappyBird-v1')
    episode, rew_sum = 20, 0
    for _ in range(episode):
        env.reset()
        info = None
        while True:
            action, _, __ = flappy_bird_rule(None if info is None else info['origin'], FlappyBirdEnv.rand_action())
            observation, reward, done, info = env.step(action)
            cv2.imshow('obs', cv2.merge(observation))
            if done:
                pygame.quit()
                print('score: %d' % env.score.score)
                rew_sum += env.score.score
                break
            env.render()
            time.sleep(0.1)
    print('avg: %f' % (rew_sum / episode))
